# _*_ coding: utf-8 _*_
import pickle
from random import sample, shuffle, uniform, randint

import numpy as np
from tqdm import trange

from Encodings import MultiToursSolution


def copy(target):
    return pickle.loads(pickle.dumps(target))


class VNS:
    def __init__(self, data):
        self.name = __class__.__name__
        # -- read instance from data --
        self.data = copy(data)
        self.V = self.data.cities  # city list
        self.n = self.data.n  # number of cities
        self.m = self.data.m  # number of salesmen
        self.A = self.data.A  # accessibility matrix
        self.Depot = self.data.depot  # depot
        self.V.remove(self.Depot)  # exclude depot
        self.Distance = self.data.distance_matrix  # distance matrix
        self.tour_cost = self.data.cost  # cost function for a single tour
        self.access = self.data.access  # fast computation access func
        # -- parameters ---
        self.s_max = self.n  # size of neighborhood
        self.ita = 0.1  # removing probability
        self.opt_depth = 1  # 2-opt maximal depth
        # -- solution initialization --
        self.solution = MultiToursSolution().init_solution(cities=self.V, distance=self.Distance, access=self.access,
                                                           depot=self.Depot, m=self.m)
        self.optima = self.solution
        self.evo_list = []
    
    def optimize(self, iteration=1000, trial=1):
        pbar = trange(iteration)
        pbar.set_description(f'{self.name} on {self.data.name} trial {trial}')
        for _ in pbar:
            new_solution = self.shaking(x=self.solution)
            new_solution = self.local_search(x=new_solution)
            if new_solution.total < self.solution.total:
                self.solution = copy(new_solution)
            self.evo_list.append(self.solution.total)
            pbar.set_postfix_str(f'best: {self.solution.total:.2f}')
        self.optima = self.solution
        return self
    
    def shaking(self, x):
        neighbors = [self.interchange(x) if uniform(0, 1) < 0.5 else self.relocate(x) for _ in range(self.s_max)]
        fit = np.array([1.0 / (1 + x.total) for x in neighbors])  # fitness = 1 / (1 + length)
        fit = fit / fit.sum()  # fitness proportion
        winner = np.random.choice(a=np.arange(0, len(neighbors)), size=1, p=fit)[0]  # roulette wheel selection
        return neighbors[winner]
    
    def interchange(self, x):
        y = copy(x)
        k = randint(0, self.m - 1)
        i = randint(1, len(y.tours[k]) - 2)
        j = randint(1, len(y.tours[k]) - 2)
        y.tours[k][i], y.tours[k][j] = y.tours[k][j], y.tours[k][i]
        return y.calc_cost(cost_func=self.tour_cost)
    
    def relocate(self, x):
        y = copy(x)
        k = randint(0, self.m - 1)
        i = randint(1, len(y.tours[k]) - 2)
        city = y.tours[k][i]
        y.tours[k].pop(i)  # remove i
        k = sample(self.access(city=city), 1)[0]
        i = randint(1, len(y.tours[k]) - 1)
        y.tours[k].insert(i, city)
        return y.calc_cost(cost_func=self.tour_cost)
    
    def local_search(self, x):
        remove = []
        for tour in x.tours:
            for i in range(len(tour) - 2, 0, -1):
                if uniform(0, 1) < self.ita:
                    remove.append(tour[i])
                    tour.pop(i)
        shuffle(remove)
        for city in remove:
            x.optimal_insert(city=city, distance=self.Distance, access=self.access)
        x.two_opt(distance=self.Distance, depth=self.opt_depth)
        return x.calc_cost(cost_func=self.tour_cost)
